package LearnDataStructure.c_树结构.有序树.二叉树.平衡二叉树;

import LearnDataStructure.c_树结构.BSTNode;
import LearnDataStructure.c_树结构.MyBinarySearchTreeAPI;
import LearnDataStructure.c_树结构.有序树.二叉树.b_MyBinaryTreebyLinked;

import java.util.Comparator;

/**
 * @version: 1.0
 * @Author: wuqiu
 * @date: 2023-07-09 17:15
 * @description: LeetcodeLearnAlgorithm -> LearnDataStructure.c_树结构.有序树.二叉树.平衡二叉树
 */
public class c_MyAVLTree<K, V> extends b_MyBinaryTreebyLinked<K, V> implements MyBinarySearchTreeAPI<K, V> {

    public c_MyAVLTree(Comparator<K> comparator) {
        super(comparator);
    }

    public c_MyAVLTree() {
    }

    @Override
    public BSTNode<K, V> insert(K key, V value) {
        BSTNode<K, V> newNode = super.insert(key, value);
        BSTNode<K, V>[] pqs = findFirstUnbalance(newNode);
        if (pqs != null) {
            reBalance(pqs);
        }
        return newNode;
    }

    /**
     * 寻找node顺着树向上
     * 第一个不平衡的节点
     * 如果有，一定是祖父级别及其以上的节点；因为这个是理论的出来的
     *
     * @param node
     * @return
     */
    private BSTNode<K, V>[] findFirstUnbalance(BSTNode<K, V> node) {
        if (node == null) {
            return null;
        }
        BSTNode<K, V> son = node;
        BSTNode<K, V> parent = node.parent;
        if (parent == null) {
            return null;
        }
        BSTNode<K, V> grand = parent.parent;
        if (grand == null) {
            return null;
        }
        if (unBalance(grand)) {
            return new BSTNode[]{grand, parent, son};//找到了，返回
        } else {
            return findFirstUnbalance(parent);//递归顺着树向上接着找
        }
    }

    /**
     * 四种类型
     * 进行旋转的操作
     * 是平衡操作的，核心代码逻辑
     *
     * @param pqs
     */
    private void reBalance(BSTNode<K, V>[] pqs) {
        if (pqs == null) {
            return;
        }
        BSTNode<K, V> p = pqs[0];
        BSTNode<K, V> q = pqs[1];
        BSTNode<K, V> s = pqs[2];
        if (q.isRight() && s.isRight()) {//右右型，以p为中心逆时针旋转
            leftRotate(p, q);
        } else if (q.isLeft() && s.isLeft()) {//左左型，以p为中心顺时针旋转
            rightRotate(p, q);
        } else if (q.isLeft() && s.isRight()) {//左右型
            //q，s左旋，变为左左型
            leftRotate(q, s);
            rightRotate(p, s);
        } else {//右左型
            rightRotate(q, s);
            leftRotate(p, s);
        }
    }

    @Override
    public void remove(K key) {
        BSTNode<K, V> node = getValue(key);
        if (node == null) {
            return;
        }
        BSTNode<K, V> parent = node.parent;
        BSTNode<K, V> left = node.left;
        BSTNode<K, V> right = node.right;
        if (left == null && right == null) {        //node叶子节点
            removeNode(node);
            needBalance(parent);//判断父节点是否平衡
        } else if (right == null) {                 //node仅有左孩子
            BSTNode<K, V> pre = max(left);
            BSTNode<K, V> parnetOfpre = pre.parent;
            K keyOfpre = pre.key;
            V valueOfpre = pre.value;
            removeNode(pre);
            node.key = keyOfpre;
            node.value = valueOfpre;
            needBalance(parnetOfpre);//判断父节点是否平衡
        } else if (left == null) {                  //node仅有右孩子
            BSTNode<K, V> suc = min(right);
            BSTNode<K, V> parentOfsuc = suc.parent;
            K keyOfsuc = suc.key;
            V valueOfsuc = suc.value;
            removeNode(suc);
            node.key = keyOfsuc;
            node.value = valueOfsuc;
            needBalance(parentOfsuc);//判断父节点是否平衡
        } else {                                    //node有两个孩子
            BSTNode<K, V> suc = min(right);
            BSTNode<K, V> parentOfsuc = suc.parent;
            removeNode(node);
            needBalance(parentOfsuc);//判断父节点是否平衡
        }
        size--;
    }

    /**
     * 检查当前节点是否平衡
     * 然后逐渐检查全树是否平衡
     * @param node
     */
    public void needBalance(BSTNode<K, V> node) {
        if (node == null) {//递归出口
            return;
        }
        BSTNode<K, V> right = node.right;
        BSTNode<K, V> left = node.left;
        int difference = getHeight(node.right) - getHeight(node.left);//有正负
        if (difference > 1) {//node右侧有不平衡
            BSTNode<K, V> p = node;
            BSTNode<K, V> q = right;

            //这里需要判断非右右型左左型的情况
            if (getHeight(right.right) >= getHeight(right.left)) {//右右型
                leftRotate(p, q);
            } else {//右左型
                BSTNode<K, V> s = right.left;
                rightRotate(q,s);
                leftRotate(p,s);
            }

            needBalance(right);
        } else if (difference < -1) {//node左侧有不平衡
            BSTNode<K, V> p = node;
            BSTNode<K, V> q = left;

            //这里需要判断非右右型左左型的情况
            if (getHeight(left.left) >= getHeight(left.right)) {//左左型
                rightRotate(node, left);
            } else {//左右型
                BSTNode<K, V> s = left.right;
                leftRotate(q,s);
                rightRotate(p,s);
            }
            needBalance(left);
        } else {//node为节点的子树平衡，顺着树向上检查
            needBalance(node.parent);
        }
    }
}